from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms



class Mydata(Dataset):
    def __init__(self, root_dir, label_dir, transform=None):  # 新增 transform 参数
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.image_path = os.listdir(self.path)
        self.transform = transform  # 保存 transform 用于后续处理

    def __getitem__(self, idx):
        image_name = self.image_path[idx]
        image_item_path = os.path.join(self.path, image_name)
        image = Image.open(image_item_path)
        label = self.label_dir

        # 关键：在数据集内部对单个图像应用 transform（转为 Tensor）
        if self.transform:
            image = self.transform(image)

        return image, label  # 此时 image 已经是 Tensor

    def __len__(self):
        return len(self.image_path)


# 定义 transform（Resize + 转 Tensor）
dataset_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

train_root_dir = "train"
ants_label_dir = "ants"
# 创建数据集时传入 transform，确保返回的是 Tensor
ants_dataset = Mydata(train_root_dir, ants_label_dir, transform=dataset_transform)
# DataLoader 加载的 batch 中，图像已经是 Tensor
test_loader = DataLoader(dataset=ants_dataset, batch_size=16, shuffle=True, num_workers=0, drop_last=False)


writer = SummaryWriter("loader_logs")
step = 0
for data in test_loader:
    ants_imgs, ants_labels = data  # ants_imgs 是 shape 为 [16, 3, 256, 256] 的 Tensor（批量图像）
    writer.add_images("test",ants_imgs,step)
    step += 1
writer.close()